import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from .feature_decorator import support_predict

@support_predict
def prepare_features(df):
    """增强的季节性特征"""
    dates = pd.to_datetime(df.iloc[:, 0])
    year = dates.dt.year
    month = dates.dt.month
    
    # 细化的周期性特征
    month_sin = np.sin(2 * np.pi * month/12)
    month_cos = np.cos(2 * np.pi * month/12)
    half_year_sin = np.sin(4 * np.pi * month/12)
    half_year_cos = np.cos(4 * np.pi * month/12)
    quarter_sin = np.sin(8 * np.pi * month/12)
    quarter_cos = np.cos(8 * np.pi * month/12)
    
    # 季节特征
    def get_season(month):
        if 3 <= month <= 5:
            return 1    # 春季
        elif 6 <= month <= 8:
            return 2    # 夏季
        elif 9 <= month <= 11:
            return 3    # 秋季
        else:
            return 4    # 冬季
    
    season = np.array([get_season(m) for m in month])
    
    # 用电高峰期标记
    peak_season = np.where(
        ((month >= 6) & (month <= 8)) |  # 夏季
        ((month >= 12) | (month <= 2)),  # 冬季
        1, 0
    )
    
    # 年度趋势特征
    years_since_start = year - year.min()
    
    # 月度用电特征独热编码
    month_columns = [f'month_{i}' for i in range(1, 13)]
    month_dummies = pd.DataFrame(0, index=range(len(df)), columns=month_columns)
    for i, m in enumerate(month, 1):
        month_dummies.loc[i-1, f'month_{m}'] = 1
    
    X = np.column_stack([
        year,
        month,
        month_sin,
        month_cos,
        half_year_sin,
        half_year_cos,
        quarter_sin,
        quarter_cos,
        season,
        peak_season,
        years_since_start,
        month_dummies
    ])
    
    # 标准化（排除月份独热编码部分）
    scaler = StandardScaler()
    X[:, :11] = scaler.fit_transform(X[:, :11])
    y = df.iloc[:, 1].values
    return X, y